import sqlite3
import pandas as pd

class SQLiteTool:
    def __init__(self, db_name):
        self.conn = sqlite3.connect(db_name)
        self.cursor = self.conn.cursor()

    def __del__(self):
        self.cursor.close()
        self.conn.close()

    def query_last_row_id(self):
        return self.execute_one("SELECT last_insert_rowid();")

    def execute_query(self, query, params=None):
        try:
            if params:
                self.cursor.execute(query, params)
            else:
                self.cursor.execute(query)
            self.conn.commit()
            return self.cursor.fetchall()
        except sqlite3.Error as e:
            print(f"执行查询时出错: {e}")
            self.conn.rollback()
            raise Exception("执行查询时出错", e)

    def insert(self, table_name, data):
        columns = ", ".join(data.keys())
        placeholders = ", ".join(["?"] * len(data))
        query = (
            f"INSERT OR IGNORE INTO {table_name} ({columns}) VALUES ({placeholders})"
        )
        # print("sql = ", query)
        # print("data = ", tuple(data.values()))
        self.execute_query(query, tuple(data.values()))
        return self.execute_query("SELECT last_insert_rowid();")

    def insert_many(self, table_name, data_list):
        if not data_list:
            return
        columns = ", ".join(data_list[0].keys())
        placeholders = ", ".join(["?"] * len(data_list[0]))
        query = (
            f"INSERT OR IGNORE INTO {table_name} ({columns}) VALUES ({placeholders})"
        )
        params = [tuple(d.values()) for d in data_list]
        try:
            self.cursor.executemany(query, params)
            self.conn.commit()
            return self.execute_query("SELECT last_insert_rowid();")
        except sqlite3.Error as e:
            print(f"批量插入时出错: {e}")
            self.conn.rollback()
            return -1

    def update(self, table_name, data, condition):
        set_clause = ", ".join([f"{key} =?" for key in data.keys()])
        query = f"UPDATE {table_name} SET {set_clause} WHERE {condition}"
        return self.execute_query(query, tuple(data.values()))

    def update_many(self, table_name, data_list, condition_template):
        if not data_list:
            return
        set_clause = ", ".join([f"{key} =?" for key in data_list[0].keys()])
        query = f"UPDATE {table_name} SET {set_clause} WHERE {condition_template}"
        params = []
        for data in data_list:
            param = tuple(data.values())
            params.append(param)
        try:
            self.cursor.executemany(query, params)
            self.conn.commit()
            return True
        except sqlite3.Error as e:
            print(f"批量更新时出错: {e}")
            self.conn.rollback()
            return False

    def delete(self, table_name, condition):
        query = f"DELETE FROM {table_name} WHERE {condition}"
        return self.execute_query(query)

    def delete_many(self, table_name, conditions):
        if not conditions:
            return
        query = f"DELETE FROM {table_name} WHERE?"
        try:
            self.cursor.executemany(query, [(condition,) for condition in conditions])
            self.conn.commit()
            return True
        except sqlite3.Error as e:
            print(f"批量删除时出错: {e}")
            self.conn.rollback()
            return False

    def select(self, table_name, columns="*", condition=None, order_by=None, limit=None):
        query = f"SELECT {columns} FROM {table_name}"
        if condition:
            query += f" WHERE {condition}"
        if order_by:
            query += f" ORDER BY {order_by}"
        if limit:
            query += f" LIMIT {limit}"
        return self.execute_query(query)

    def select_list_to_dataframe(self, table_name, columns="*", condition=None, order_by=None, limit=None):
        # SQL查询语句
        data = self.select(table_name, columns, condition, order_by, limit)
        # 下面为将获取的数据转化为dataframe格式
        # 获取表的列名
        column_names = [description[0] for description in self.cursor.description]
        # 将查询结果转换为 Pandas DataFrame
        df = pd.DataFrame(data, columns=column_names)
        return df

if __name__ == "__main__":
    db_tool = SQLiteTool("financial_data.db")
    # 创建一个示例表
    db_tool.execute_query(
        """CREATE TABLE IF NOT EXISTS t_users 
                             (id INTEGER PRIMARY KEY AUTOINCREMENT, 
                              name TEXT, 
                              age INTEGER)"""
    )
    # 插入单条数据
    db_tool.insert("t_users", {"name": "Alice", "age": 25})
    # 批量插入数据
    data_list = [{"name": "Bob", "age": 30}, {"name": "Charlie", "age": 35}]
    db_tool.insert_many("t_users", data_list)
    # 查询所有数据
    result = db_tool.select("t_users")
    for row in result:
        print(row)
    # 更新单条数据
    db_tool.update("t_users", {"age": 26}, "name = 'Alice'")
    # 批量更新数据
    update_list = [{"age": 31}, {"age": 36}]
    db_tool.update_many("t_users", update_list, "name = 'Alice'")
    # 删除单条数据
    db_tool.delete("t_users", "name = 'Bob'")
    # 批量删除数据
    db_tool.delete_many("t_users", ["name = 'Charlie'"])
